# 弹簧阻尼系统的LQR控制:跟踪问题：xd不再是一个常数。如xd = [0;0.2]
# 质量块位置跟随sinx

from datetime import datetime
import F_LQR_Gain 
import numpy as np
import math
import matplotlib.pyplot as plt
import pandas as pd
import control as ct
import time
import scipy.linalg as la
from scipy.signal import StateSpace
import F3_InputAugmentMatrix_Delta_U

# 1.定义系统
m = 1
k = 1
b = 0.5
# x[k+1] = Ax[k] + Bu[k]
A = np.array([[0,1],[-k/m,-b/m]])
B = np.array([[0],[1/m]])
C = np.array([[1,0]])
D = 0
# 创建状态空间模型
# sys = StateSpace(A, B, C, D)
Ts = 0.1
sys_continuous = ct.StateSpace(A, B, C, D)
sys_discrete = ct.sample_system(sys_continuous, Ts, method='zoh')
# 获取行数
n = len(A)
# 获取列数
p =  B.shape[1]
# 离散化后的状态矩阵和输入矩阵
A = sys_discrete.A
B = sys_discrete.B
# 目标:
w = 0.5*math.pi
xd = np.array([[0.1],[0]]).reshape(1,-1).T
AD = ct.sample_system(ct.StateSpace(np.array([[0,1],[-w*w,0]]),
                                    np.array([[0],[0]]),np.array([[0,0]]),0),Ts,method='zoh').A
# 2.初始化系统
x0 = np.array([[0],[0]]).reshape(1,-1)
x = x0
# 输入
print("当前时间是:",datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
u0 = np.array([[0]]).reshape(1,-1)
u = u0
xa = np.hstack((x,xd.T,u))
# 定义系统运行步数
k_steps = 200
x_history = np.zeros((n,k_steps+1))
x_history[:,0] = x
u_history = np.zeros((p,k_steps))
u_history[:,0] = u
xa_history = np.zeros((2*n+p,k_steps+1))
xa_history[:,0] = xa.T.flatten()
# 设置权重
Q = np.array([[1,0],[0,1]])
S = np.array([[1,0],[0,1]])
R=0.1

N = k_steps
P_k = S
# 组合新的矩阵
Aa,Ba,Qa,Sa,R= F3_InputAugmentMatrix_Delta_U.InputAugmentMatrix_Delta_U(A,B,Q,R,S,AD)
# 3.计算反馈增益
F = F_LQR_Gain.LQR_Gain(Aa,Ba,Qa,R,Sa)

for k in range(1,k_steps+1):
    if k == 50 :
        xd = np.vstack((xd[0,0],-0.2))        
    elif k == 100 :
        xd = np.vstack((xd[0,0],0.2)) 
    elif k == 150 :
        xd = np.vstack((xd[0,0],-0.2)) 
    elif k == 200 :
        xd = np.vstack((xd[0,0],0.2)) 

    Delta_u = -F@xa_history[:,k-1].reshape(-1,1)
    u = Delta_u + u
    x=A@x_history[:,k-1].reshape(-1,1)+B@u
    xd = AD@xd
    xa = np.vstack((x,xd,u))
    xa_history[:,k]=xa.T
    x_history[:,k]=x.T
    u_history[:,k-1]=u
 
"""
# # 写入数据到表格中
from openpyxl import Workbook
# 创建 Excel 工作簿
wb = Workbook()
ws = wb.active
# 将矩阵数据写入工作表
for row in xa_history:
    ws.append(row.tolist())
# 保存 Excel 文件
wb.save('xa_history.xlsx')
wb = Workbook()
ws = wb.active

# 将矩阵数据写入工作表
for row in u_history:
    ws.append(row.tolist())
# 保存 Excel 文件
wb.save('u_history.xlsx')

"""

# 4.制图
plt.figure()
plt.subplot(3,1,1)
plt.title('Input and State Plot')
plt.xlabel('steps')
plt.ylabel('State')
plt.plot(x_history[0,:],label=f'State 1')
plt.plot(xa_history[2,:],label=f'desired 1')

plt.legend()
plt.subplot(3,1,2)
plt.title('Input and State Plot')
plt.xlabel('steps')
plt.ylabel('State')
plt.plot(x_history[1,:],label=f'State 2')
plt.plot(xa_history[3,:],label=f'desired 2')
plt.legend()

plt.subplot(3,1,3)
plt.ylabel('Input')
for i in range(0,p):
    plt.plot(u_history[i,:],label=f'Input {i+1}')
    
plt.legend()
plt.grid()
plt.show()